"""
EEC 284 Assignment 3 Problem 2
TFO Signal Analysis
Author: YOUR NAME
Instructions:
  - Complete all "TODO" sections
  - Do not modify function signatures
  - Submit this file with all functions completed and working
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import scipy
from scipy.signal import butter, filtfilt, spectrogram

# ------------------ Local Functions ------------------
def plot_time(signal, Fs, start, stop):
    t = np.arange(stop - start)/Fs
    s = signal[start:stop]
    plt.plot(t,s)
    plt.title("Time vs Amplitude")
    plt.xlabel("Time (s)")
    plt.ylabel("Signal")
    plt.show()

def plot_fft(signal, Fs, title="FFT Magnitude Spectrum", f_max=None):
    """
    Calculates and plots the FFT of a signal in dB.

    Args:
        signal (np.ndarray): The input 1D signal.
        Fs (int): The sampling frequency in Hz.
        title (str, optional): The title for the plot. 
                               Defaults to "FFT Magnitude Spectrum".
        f_max (float, optional): The maximum frequency for the x-axis. 
                                 Defaults to the Nyquist frequency (Fs / 2).
    """
    eps = 1e-12
    n = signal.size
    yf = scipy.fft.fft(signal)/n
    xf = scipy.fft.fftfreq(n, 1 / Fs)
    
    pos_freqs_mask = np.where(xf >= 0)
    xf_pos = xf[pos_freqs_mask]
    yf_mag_pos = 20 * np.log10(np.abs(yf[pos_freqs_mask])+eps)

    if f_max is None:
        f_max = Fs / 2

    plt.figure(figsize=(12, 6))
    plt.plot(xf_pos, yf_mag_pos)
    plt.title(title) 
    plt.xlabel("Frequency [Hz]")
    plt.ylabel("Magnitude [dB]")
    plt.grid(True)
    plt.xlim(0, f_max) 
    plt.show()

def estimate_carrier_frequencies(signal, Fs):
    """
    TODO: Use FFT to find two dominant carrier frequencies
    Inputs:
        signal - 1D signal from any detector
        Fs - sampling rate in Hz
    Output:
        f1 - estimated carrier frequency 1
        f2 - estimated carrier frequency 2
    """
    yf = scipy.fft.fft(signal)
    n = signal.size
    step = 1/Fs
    xf = scipy.fft.fftfreq(n, step)
    pos_freqs = np.where(xf>0)
    xf_freq_pos = xf[pos_freqs]
    yf_mag_pos = np.abs(yf[pos_freqs])
    peaks, properties = scipy.signal.find_peaks(yf_mag_pos, np.max(yf_mag_pos)*0.2, distance = 1000)
    peak_heights = properties['peak_heights']
    peak_freqs = xf_freq_pos[peaks]
    peak_freqs = [peak for peak in peak_freqs if peak > 10]
    return peak_freqs




def iq_demodulate(signal, fc, Fs):
    """
    TODO: Perform I/Q demodulation of input signal
    Inputs:
        signal - 1D signal to demodulate
        fc - carrier frequency to demodulate
        Fs - sampling rate
    Output:
        demod - demodulated signal envelope via sqrt(I^2 + Q^2)
    """
    pass
    step = 1/Fs
    n = signal.size
    t = np.arange(n)/Fs
    I = np.cos(fc * 2 * np.pi * t)
    Q = np.sin(fc* 2 * np.pi * t)
    I_demod = I * signal
    Q_demod = Q * signal
    sum = np.sqrt(np.power(I_demod,2) + np.power(Q_demod,2))
    sum = lowpass_filter(sum,15,Fs)
    #plot_fft(signal,Fs,"Signal FFT",20)
    #plot_fft(sum,Fs,"Sum FFT", 20)

    return sum
    

def lowpass_filter(data, cutoff, Fs):
    """
    TODO: Apply low-pass filtering for I/Q demodulation
    Inputs:
        data   - 1D signal to demodulate
        cutoff - Cutoff frequency for low-pass in Hz (use 15Hz here)
        Fs     - Sampling rate in Hz
    Output:
        filtered - Filtered version of input
    """
    sos = scipy.signal.butter(10,cutoff,'lp',fs=Fs,output='sos')
    filtered = scipy.signal.sosfilt(sos,data)
    return filtered

def plot_spectrogram(signal, Fs, window_sec, overlap):
    """
    TODO: Plot spectrogram of input 1D signal
    Requirements:
      - Frequency on x-axis, time on y-axis
      - Spectrum type: 'power'
      - X-axis limit: [0, 5] Hz
    Inputs:
      signal      - A 1D signal
      Fs          - Sampling frequency (e.g., 80 Hz after downsampling)
      window_sec  - Window size in seconds (e.g., 60)
      overlap     - Fractional overlap (e.g., 0.5)
    """
    window_samples = int(window_sec * Fs)
    noverlap = int(overlap*window_samples)
    print(f"window_samples: {window_samples} noverlap: {noverlap}")
    f, t, Sxx = scipy.signal.spectrogram(signal, Fs, nperseg = window_samples,noverlap = noverlap, mode = 'magnitude')
    plt.pcolormesh(f, t, 10*np.log10(Sxx.T), shading='gouraud')
    plt.xlim(0,5)
    plt.ylabel('Time [sec]')
    plt.xlabel('Frequency [hz]')
    plt.show()